package cn.imerge;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import io.socket.client.Ack;
import io.socket.client.IO;
import io.socket.client.Socket;
import org.apache.commons.lang3.StringUtils;

import java.net.URI;
import java.net.URISyntaxException;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

public class ChatBot {

    public static class Conversation {
        private String id;

        private String conversationId;

        private String parentId;

        private Long lastActive;

        public String getId() {
            return id;
        }

        public void setId(String id) {
            this.id = id;
        }

        public String getConversationId() {
            return conversationId;
        }

        public void setConversationId(String conversationId) {
            this.conversationId = conversationId;
        }

        public String getParentId() {
            return parentId;
        }

        public void setParentId(String parentId) {
            this.parentId = parentId;
        }

        public Long getLastActive() {
            return lastActive;
        }

        public void setLastActive(Long lastActive) {
            this.lastActive = lastActive;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Conversation that = (Conversation) o;
            return Objects.equals(id, that.id);
        }

        @Override
        public int hashCode() {
            return Objects.hash(id);
        }
    }

    private String sessionToken;

    private String bypassNode;

    private String userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/109.0.0.0 Safari/537.36 Edg/109.0.1518.70";

    private Long tokenExpires = new Date().getTime();

    private AtomicBoolean running = new AtomicBoolean(false);

    private AtomicBoolean init = new AtomicBoolean(false);

    private Long nextCheckToken = new Date().getTime();

    private AtomicBoolean ready = new AtomicBoolean(false);

    // 缓存用户信息
    private Map<String, Conversation> conversationMap = new ConcurrentHashMap<>();

    private String auth;

    private AtomicBoolean stopCheckToken = new AtomicBoolean(false);

    private Socket socket;

    public ChatBot(String sessionToken) throws URISyntaxException {
        this(sessionToken, "https://gpt.pawan.krd");
    }

    public ChatBot(String sessionToken, String bypassNode) throws URISyntaxException {
        this.sessionToken = sessionToken;
        this.bypassNode = bypassNode;

        this.sessionToken = sessionToken;

        IO.Options options = IO.Options.builder()
                .setTransports(new String[]{"websocket", "polling"})
                .build();

        this.socket = IO.socket(new URI(bypassNode + "/?client=python&version=1.0.2&versionCode=102"), options).connect();

        this.socket.on("connect", objects -> {
            System.out.println("socket 服务开启成功");
            init.set(true);
        });

        this.socket.on("disconnect", objects -> {
            System.out.println("socket 断开连接");
        });

        this.waitForInited();

        this.running = new AtomicBoolean(true);

        // 启动token过期扫描线程
        new Thread(() -> {
            while (running.get()) {
                if (nextCheckToken > new Date().getTime()) {
                    try {
                        Thread.sleep(100);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                    continue;
                }

                // 检查token的有效期是否在120秒以内
                var offset = TimeUnit.MINUTES.toMillis(2);

                if (tokenExpires < (new Date().getTime() - offset) || StringUtils.isBlank(auth)) {
                    getToken();
                }
                nextCheckToken = new Date().getTime() + offset;
                try {
                    Thread.sleep(5000);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }).start();

    }

    public Conversation addConversation(String id) {

        var conversation = new Conversation();
        conversation.setId(id);
        conversation.setConversationId(null);
        conversation.setParentId(UUID.randomUUID().toString());
        conversation.setLastActive(new Date().getTime());

        conversationMap.put(id, conversation);

        return conversation;
    }

    public void waitForInited() {
        while (!this.init.get()) {
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    public void waitForReady() {
        while (!this.ready.get()) {
            try {
                Thread.sleep(1500);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    public Conversation getConversation(String id) {
        var conversation = conversationMap.get(id);

        if (conversation == null) {
            conversation = addConversation(id);
        } else {
            conversation.setLastActive(new Date().getTime());
        }

        return conversation;
    }

    public void getToken() {
        try {
            Thread.sleep(1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        this.socket.emit("getSession", this.sessionToken, new Ack() {
            @Override
            public void call(Object... objects) {
                System.out.println("获取到更新后的token");

                System.out.println(objects.length);

                JSONObject tokenJson = (JSONObject) JSON.toJSON(objects[0]);

                System.out.println(tokenJson);

                if (tokenJson.containsKey("error")) {
                    System.out.println("更新token失败，服务器返回错误: " + tokenJson.getString("error"));
                    return;
                }

                sessionToken = tokenJson.getString("sessionToken");

                auth = tokenJson.getString("auth");
                // 2023-01-31T15:16:58.065Z
                try {
                    tokenExpires = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSSZ", Locale.US).parse(tokenJson.getString("expires").replace("Z","+0000")).getTime();
                } catch (ParseException e) {
                    e.printStackTrace();
                }
                ready.set(true);
            }
        });

    }

    public String ask(String question, String id) throws InterruptedException {
        if (StringUtils.isBlank(this.auth) || validateToken()) {
            getToken();
        }

        var conversation = this.getConversation(id);

        JSONObject req = new JSONObject();

        req.put("auth", this.auth);
        req.put("prompt", question);
        req.put("parentId", conversation.getParentId());
        req.put("conversationId", conversation.getConversationId());

        CountDownLatch waitResponse = new CountDownLatch(1);

        final String[] responseMessage = {null};

        this.socket.emit("askQuestion", req, new Ack() {
            @Override
            public void call(Object... args) {
                System.out.println("获取响应 : " + JSON.toJSONString(args));

                JSONObject response = JSON.parseObject(args[0].toString());

                if (response.containsKey("error")){
                    return;
                }
                responseMessage[0] = response.getString("answer");

                waitResponse.countDown();
            }
        });
        waitResponse.await();
        return responseMessage[0] != null ? responseMessage[0] : null;
    }

    public boolean validateToken() {
        if (StringUtils.isBlank(this.auth)) {
            return false;
        }

        System.out.println(this.auth.split("\\.")[1]);

        JSONObject tokenParams = JSON.parseObject(new String(Base64.getDecoder().decode(this.auth.split("\\.")[1])));

        System.out.println(tokenParams);

        return new Date().getTime() <= tokenParams.getLong("exp");
    }

    public void close() {
        this.socket.close();
        this.running = new AtomicBoolean(false);
    }


}
